package com.yololabelmaker.view;

import com.yololabelmaker.CImgArea;
import com.yololabelmaker.LabelRect;
import com.yololabelmaker.YOLOLabelLoader;

import javax.swing.*;
import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.event.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import javax.swing.event.ListSelectionEvent;
import javax.swing.event.ListSelectionListener;

public class DatasetBrowserPanel extends JPanel {
    private CImgArea imgArea;
    private File datasetDir;
    private DefaultListModel<String> imageListModel;
    private JList<String> imageList;
    private JLabel imagePreviewLabel;
    private JCheckBox showLabelsCheckBox;
    private JLabel annotationCountLabel;
    private List<LabelRect> currentImageLabels;

    public DatasetBrowserPanel(CImgArea imgArea, File datasetDir) {
        this.imgArea = imgArea;
        this.datasetDir = datasetDir;
        this.imageListModel = new DefaultListModel<>();
        this.currentImageLabels = new ArrayList<>();
        initUI();
        loadDatasetImages();
    }

    private void initUI() {
        setLayout(new BorderLayout());

        // 创建主面板，使用边框布局
        JPanel mainPanel = new JPanel(new BorderLayout());

        // 创建图片列表
        imageList = new JList<>(imageListModel);
        imageList.setSelectionMode(ListSelectionModel.SINGLE_SELECTION);
        imageList.setCellRenderer(new ImageListCellRenderer());
        JScrollPane listScrollPane = new JScrollPane(imageList);
        listScrollPane.setPreferredSize(new Dimension(250, 600));

        // 创建图片预览面板
        JPanel previewPanel = new JPanel(new BorderLayout());
        imagePreviewLabel = new JLabel("选择一张图片进行预览", SwingConstants.CENTER);
        imagePreviewLabel.setBorder(BorderFactory.createEmptyBorder(10, 10, 10, 10));
        JScrollPane previewScrollPane = new JScrollPane(imagePreviewLabel);
        previewPanel.add(previewScrollPane, BorderLayout.CENTER);

        // 创建信息面板
        JPanel infoPanel = new JPanel(new FlowLayout(FlowLayout.LEFT));
        annotationCountLabel = new JLabel("标注数量: 0");
        showLabelsCheckBox = new JCheckBox("显示标注");
        showLabelsCheckBox.setSelected(true);
        showLabelsCheckBox.addActionListener(e -> updateImagePreview());
        infoPanel.add(showLabelsCheckBox);
        infoPanel.add(annotationCountLabel);
        previewPanel.add(infoPanel, BorderLayout.NORTH);

        // 创建导航按钮面板
        JPanel navButtonPanel = new JPanel();
        JButton prevButton = new JButton("上一张");
        JButton nextButton = new JButton("下一张");
        JButton openButton = new JButton("在主窗口打开");

        navButtonPanel.add(prevButton);
        navButtonPanel.add(nextButton);
        navButtonPanel.add(openButton);

        previewPanel.add(navButtonPanel, BorderLayout.SOUTH);

        // 添加到主面板
        mainPanel.add(listScrollPane, BorderLayout.WEST);
        mainPanel.add(previewPanel, BorderLayout.CENTER);

        add(mainPanel, BorderLayout.CENTER);

        // 添加图片列表选择事件监听器
        imageList.addListSelectionListener(e -> {
            if (!e.getValueIsAdjusting() && !imageList.isSelectionEmpty()) {
                updateImagePreview();
            }
        });

        // 上一张按钮事件
        prevButton.addActionListener(e -> {
            int currentIndex = imageList.getSelectedIndex();
            if (currentIndex > 0) {
                imageList.setSelectedIndex(currentIndex - 1);
                imageList.ensureIndexIsVisible(currentIndex - 1);
            }
        });

        // 下一张按钮事件
        nextButton.addActionListener(e -> {
            int currentIndex = imageList.getSelectedIndex();
            if (currentIndex < imageListModel.size() - 1) {
                imageList.setSelectedIndex(currentIndex + 1);
                imageList.ensureIndexIsVisible(currentIndex + 1);
            }
        });

        // 在主窗口打开按钮事件
        openButton.addActionListener(e -> {
            if (!imageList.isSelectionEmpty()) {
                String selectedImagePath = imageList.getSelectedValue();
                try {
                    BufferedImage image = ImageIO.read(new File(selectedImagePath));
                    if (image != null) {
                        imgArea.setImage(image, selectedImagePath);
                    }
                } catch (IOException ex) {
                    JOptionPane.showMessageDialog(this, "打开图片时发生错误: " + ex.getMessage());
                }
            }
        });
    }

    private void loadDatasetImages() {
        try {
            Files.walk(Paths.get(datasetDir.getAbsolutePath()))
                .filter(Files::isRegularFile)
                .map(Path::toString)
                .filter(s -> s.toLowerCase().endsWith(".jpg") ||
                        s.toLowerCase().endsWith(".jpeg") ||
                        s.toLowerCase().endsWith(".png") ||
                        s.toLowerCase().endsWith(".gif"))
                .forEach(imageListModel::addElement);
        } catch (IOException e) {
            JOptionPane.showMessageDialog(this, "浏览数据集时发生错误: " + e.getMessage());
        }

        // 如果有图片，默认选择第一张
        if (imageListModel.size() > 0) {
            imageList.setSelectedIndex(0);
            // 触发选择事件以显示第一张图片的预览
            EventQueue.invokeLater(() -> {
                ListSelectionEvent event = new ListSelectionEvent(imageList, 0, 0, false);
                for (ListSelectionListener l : imageList.getListSelectionListeners()) {
                    l.valueChanged(event);
                }
            });
        }
    }

    private void updateImagePreview() {
        if (!imageList.isSelectionEmpty()) {
            String selectedImagePath = imageList.getSelectedValue();
            try {
                BufferedImage image = ImageIO.read(new File(selectedImagePath));
                if (image != null) {
                    // 加载当前图片的标签
                    loadCurrentImageLabels(selectedImagePath, image);

                    // 更新标注数量显示
                    annotationCountLabel.setText("标注数量: " + currentImageLabels.size());

                    // 创建适合预览的图像，可能包含标注
                    BufferedImage previewImage = createPreviewImage(image, 600, 400);
                    imagePreviewLabel.setIcon(new ImageIcon(previewImage));
                    imagePreviewLabel.setText("");
                    revalidate();
                    repaint();
                }
            } catch (IOException ex) {
                JOptionPane.showMessageDialog(this, "预览图片时发生错误: " + ex.getMessage());
            }
        }
    }

    private void loadCurrentImageLabels(String imagePath, BufferedImage image) {
        YOLOLabelLoader loader = new YOLOLabelLoader();
        List<LabelRect> loadedLabels = loader.loadLabels(imagePath);
        currentImageLabels.clear();
        if (loadedLabels != null) {
            // 对于每个加载的标签，更新其像素坐标
            for (LabelRect rect : loadedLabels) {
                try {
                    // 使用保存的YOLO坐标信息更新像素坐标
                    YOLOLabelLoader.updateLabelRectCoordinates(
                            rect, image.getWidth(), image.getHeight(),
                            rect.getYoloCenterX(), rect.getYoloCenterY(), 
                            rect.getYoloWidth(), rect.getYoloHeight());
                    currentImageLabels.add(rect);
                } catch (Exception e) {
                    System.err.println("更新标签坐标时发生错误: " + e.getMessage());
                }
            }
        }
    }

    private BufferedImage createPreviewImage(BufferedImage originalImage, int maxWidth, int maxHeight) {
        // 计算缩放比例
        double scale = Math.min((double) maxWidth / originalImage.getWidth(), (double) maxHeight / originalImage.getHeight());
        int newWidth = (int) (originalImage.getWidth() * scale);
        int newHeight = (int) (originalImage.getHeight() * scale);

        // 创建新的缓冲图像
        BufferedImage previewImage = new BufferedImage(newWidth, newHeight, BufferedImage.TYPE_INT_ARGB);
        Graphics2D g2d = previewImage.createGraphics();
        
        // 设置抗锯齿，使绘制更平滑
        g2d.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
        g2d.setRenderingHint(RenderingHints.KEY_TEXT_ANTIALIASING, RenderingHints.VALUE_TEXT_ANTIALIAS_ON);

        // 绘制缩放后的图片
        g2d.drawImage(originalImage, 0, 0, newWidth, newHeight, null);

        // 如果勾选了显示标注，则绘制标注框
        if (showLabelsCheckBox.isSelected() && !currentImageLabels.isEmpty()) {
            // 保存原始设置以便后续恢复
            Color originalColor = g2d.getColor();
            Stroke originalStroke = g2d.getStroke();
            Font originalFont = g2d.getFont();
            
            // 设置标注框的样式
            g2d.setColor(Color.RED);
            g2d.setStroke(new BasicStroke(2));

            for (LabelRect rect : currentImageLabels) {
                // 计算标注框在预览图中的位置和大小
                int x = (int) (rect.getX() * scale);
                int y = (int) (rect.getY() * scale);
                int width = (int) (rect.getWidth() * scale);
                int height = (int) (rect.getHeight() * scale);

                // 精确绘制标注框
                g2d.drawRect(x, y, width, height);

                // 计算YOLO格式的核心标注信息
                double yoloCenterX = rect.getYOLOCenterX(originalImage.getWidth());
                double yoloCenterY = rect.getYOLOCenterY(originalImage.getHeight());
                double yoloWidth = rect.getYOLOWidth(originalImage.getWidth());
                double yoloHeight = rect.getYOLOHeight(originalImage.getHeight());
                
                // 格式化YOLO信息文本
                String className = rect.getClassName();
                int classId = rect.getClassId();
                String labelInfo = String.format("%s (ID:%d)", className, classId);
                
                // 绘制文本背景
                g2d.setColor(new Color(255, 255, 0, 200)); // 半透明黄色背景
                FontMetrics metrics = g2d.getFontMetrics();
                int textWidth = metrics.stringWidth(labelInfo);
                int textHeight = metrics.getHeight();
                
                // 确保文本框不会超出图像边界
                int textX = x;
                int textY = Math.max(0, y - textHeight);
                int textBoxWidth = Math.max(textWidth + 10, 80); // 最小宽度80
                
                g2d.fillRect(textX, textY, textBoxWidth, textHeight);
                
                // 绘制标签信息文本
                g2d.setColor(Color.BLACK);
                g2d.drawString(labelInfo, textX + 5, textY + metrics.getAscent());
                
                // 绘制YOLO格式坐标信息（更小的字体）
                g2d.setFont(new Font(originalFont.getName(), Font.PLAIN, originalFont.getSize() - 2));
                String yoloInfo = String.format("xywh: %.3f,%.3f,%.3f,%.3f", 
                                               yoloCenterX, yoloCenterY, yoloWidth, yoloHeight);
                int yoloInfoY = textY + textHeight + g2d.getFontMetrics().getAscent();
                
                // 确保信息不会超出图像边界
                if (yoloInfoY < newHeight) {
                    g2d.fillRect(textX, textY + textHeight, 
                                g2d.getFontMetrics().stringWidth(yoloInfo) + 10, 
                                g2d.getFontMetrics().getHeight());
                    g2d.drawString(yoloInfo, textX + 5, yoloInfoY);
                }
            }
            
            // 恢复原始设置
            g2d.setColor(originalColor);
            g2d.setStroke(originalStroke);
            g2d.setFont(originalFont);
        }

        g2d.dispose();
        return previewImage;
    }

    // 自定义图片列表单元格渲染器
    private class ImageListCellRenderer extends DefaultListCellRenderer {
        @Override
        public Component getListCellRendererComponent(JList<?> list, Object value, int index, boolean isSelected, boolean cellHasFocus) {
            JLabel label = (JLabel) super.getListCellRendererComponent(list, value, index, isSelected, cellHasFocus);

            // 显示文件名而不是完整路径
            if (value instanceof String) {
                String fullPath = (String) value;
                File file = new File(fullPath);
                label.setText(file.getName());

                // 尝试加载小缩略图
                try {
                    BufferedImage image = ImageIO.read(file);
                    if (image != null) {
                        Image scaledImage = image.getScaledInstance(60, 40, Image.SCALE_SMOOTH);
                        label.setIcon(new ImageIcon(scaledImage));
                    }
                } catch (IOException e) {
                    // 如果无法加载图片，则不显示图标
                }

                // 尝试加载标签文件，显示标签数量
                YOLOLabelLoader loader = new YOLOLabelLoader();
                List<LabelRect> labels = loader.loadLabels(fullPath);
                int labelCount = (labels != null) ? labels.size() : 0;
                label.setToolTipText("标签数量: " + labelCount + "\n" + fullPath);
            }

            return label;
        }
    }
}